#pragma once
#include "arch.h"
#include "common.h"
#include "macros.h"
#include "kernel_operator.h"
#include <type_traits>
template<
    typename DstType,
    typename SrcType
>
__aicore__ inline __attribute__((always_inline)) auto getQuantMode() {
    if (std::is_same<DstType, SrcType>::value) {
        return QuantMode_t::NoQuant;
    }
    static_assert("not implement quantMode!");
    return QuantMode_t::NoQuant;
}

template<>
__aicore__ inline __attribute__((always_inline)) auto getQuantMode<half, float>() {
    return QuantMode_t::F322F16;
}

template<>
__aicore__ inline __attribute__((always_inline)) auto getQuantMode<bfloat16_t, float>() {
    return QuantMode_t::F322BF16;
}

template<
    typename Element
>
__aicore__ inline __attribute__((always_inline)) void gm2l1_Nd2Nz(
    AscendC::LocalTensor<Element> dstTensor,
    AscendC::GlobalTensor<Element> srcTensor,
    uint32_t mActual,
    uint32_t nActual,
    uint32_t srcStride,//strideA
    uint32_t mL1Size,
    uint32_t nL1Size
) {
    const uint32_t nAlignNumElement = BYTE_PER_C0 / sizeof(Element);
    //我的理解是，一个机器有BYTE_PER_C0这么多B，而一个元素的B可能比较少，所以操作tensor的时候得对其i*nAlignNumElement
    if (srcStride < STRIDE_LIMIT /*65536*/) {
        AscendC::DataCopy(
            dstTensor,
            srcTensor,
            AscendC::Nd2NzParams(
                1, 
                mActual,
                nActual,
                0,
                srcStride,
                mL1Size,
                1,
                1
            )
            /*
            struct Nd2NzParams {
                uint16_t ndNum = 0; //传输nd矩阵的数目，取值范围：ndNum∈[0, 4095]。Nd2NzParams的ndNum仅支持设置为1。
                uint16_t nValue = 0; nd矩阵的行数，取值范围：nValue∈[0, 16384]。
                uint16_t dValue = 0; nd矩阵的列数，取值范围：nValue∈[0, 16384]。
                uint16_t srcNdMatrixStride = 0; 源操作数相邻nd矩阵起始地址间的偏移，取值范围：srcNdMatrixStride∈[0, 65535]，单位：element。
                uint16_t srcDValue = 0; 源操作数同一nd矩阵的相邻行起始地址间的偏移，取值范围：srcDValue∈[1, 65535]，单位：element。
                uint16_t dstNzC0Stride = 0; ND转换到NZ格式后，源操作数中的一行会转换为目的操作数的多行。dstNzC0Stride表示，目的nz矩阵中，来自源操作数同一行的多行数据相邻行起始地址间的偏移，取值范围：dstNzC0Stride∈[1, 16384]，单位：C0_SIZE（32B）。
                uint16_t dstNzNStride = 0; 目的nz矩阵中，Z型矩阵相邻行起始地址之间的偏移。取值范围：dstNzNStride∈[1, 16384]，单位：C0_SIZE（32B）。
                uint16_t dstNzMatrixStride = 0; 目的nz矩阵中，相邻nz矩阵起始地址间的偏移，取值范围：dstNzMatrixStride∈[1, 65535]，单位：element。
            };

            通常来说，ndNum和srcNdMatrixStride是用来划分这一坨srcTensor的，把它中的某些部分化为ndNum个小的，作为原输入的ND矩阵，
            如果ndNum为1的话，srcNdMatrixStride就为0即可。因为没有第二个ND了。
            srcDValue表示一行的所含元素个数，也就是srcTensor的一行的元素个数。
            dstNzC0Stride:(简单来说就是划分完后的行数，单位是32B)


            */
        );
    } else {
        for (uint32_t i = 0; i < mActual; i++) {
            AscendC::DataCopy(
                dstTensor[i * nAlignNumElement],
                srcTensor[i * srcStride], //源ND矩阵一行有srcStride个元素.
                AscendC::Nd2NzParams(
                    1, //只有一个ND矩阵。
                    1, //这个ND矩阵只有一行。因为每次只传入srcTensor的第i行。（srcStride个元素）
                    nActual, //源矩阵的列数
                    0,      //没有第二个ND
                    srcStride, //一行的数量
                    mL1Size,
                    0,
                    1
                )
            );
        }
    }
}

template<
    typename Element
>
__aicore__ inline __attribute__((always_inline)) void l12l0_Nz2Zz(
    AscendC::LocalTensor<Element> dstTensor,
    AscendC::LocalTensor<Element> srcTensor,
    uint32_t mRound,
    uint32_t nRound,
    uint32_t srcStride,
    uint32_t ML0AlignNumElement,
    uint32_t NL0AlignNumElement
) {
    for (uint32_t i = 0; i < mRound / ML0AlignNumElement; i++) {
        AscendC::LoadData(
            dstTensor[i * nRound * ML0AlignNumElement],
            srcTensor[i * ML0AlignNumElement * NL0AlignNumElement],
            AscendC::LoadData2dParams(
                0,
                static_cast<uint16_t>(nRound / NL0AlignNumElement),
                srcStride / ML0AlignNumElement,
                0,
                0,
                false,
                0
            )
        );
    }
}

template<
    typename Element
>
__aicore__ inline __attribute__((always_inline)) void l12l0_Nz2Nz(
    AscendC::LocalTensor<Element> dstTensor,
    AscendC::LocalTensor<Element> srcTensor,
    uint32_t mRound,
    uint32_t nRound,
    uint32_t srcStride,
    uint32_t ML0AlignNumElement,
    uint32_t NL0AlignNumElement
) {
    for (uint32_t i = 0; i < nRound / NL0AlignNumElement; i++) {
        AscendC::LoadData(
            dstTensor[i * mRound * NL0AlignNumElement],
            srcTensor[i * srcStride * NL0AlignNumElement],
            AscendC::LoadData2dParams(
                0,
                static_cast<uint16_t>(mRound / ML0AlignNumElement),
                1,
                0,
                0,
                false,
                0
            )
        );
    }
}

template<
    typename ElementC,
    typename ComputeElementC
>
__aicore__ inline __attribute__((always_inline)) void l0c2gm_Nz2Nd(
    AscendC::GlobalTensor<ElementC> dstTensor,
    AscendC::LocalTensor<ComputeElementC> srcTensor,
    uint32_t mActual,
    uint32_t nActual,
    uint32_t mRound,
    uint32_t nRound,
    uint32_t dstStride
) {
    AscendC::DataCopy(
        dstTensor, 
        srcTensor,
        AscendC::DataCopyCO12DstParams(
            nActual,
            mActual,
            dstStride,
            mRound,
            getQuantMode<ElementC, ComputeElementC>(),
            0,
            false,
            true
        )
    );
}

template<
    typename ArchTag,
    typename ElementA,
    typename ElementB,
    typename ComputeElementC
>
__aicore__ inline __attribute__((always_inline)) void matmulOnL1_sync(
    AscendC::LocalTensor<ElementA> l1TensorA,
    AscendC::LocalTensor<ElementB> l1TensorB,
    AscendC::LocalTensor<ComputeElementC> l0cTensorC,
    uint32_t mL1Actual,
    uint32_t nL1Actual,
    uint32_t kL1Actual,
    uint32_t mL1Round,
    uint32_t nL1Round,
    uint32_t kL1Round,
    bool isFirstLoop
) {
    static constexpr uint32_t L0A_SIZE = ArchTag::L0ASize;
    static constexpr uint32_t L0B_SIZE = ArchTag::L0BSize;

    static constexpr uint32_t ML0AlignNumElement = C0_NUM_PER_FRACTAL;
    static constexpr uint32_t NL0AlignNumElement = C0_NUM_PER_FRACTAL;
    static constexpr uint32_t KL0AlignNumElement = BYTE_PER_C0 / sizeof(ElementA);

    // 默认L0A,L0B划分出两片空间
    static constexpr uint32_t L0A_PINGPONG_BUF_LEN = (L0A_SIZE / 2) / sizeof(ElementA);
    static constexpr uint32_t L0B_PINGPONG_BUF_LEN = (L0A_SIZE / 2) / sizeof(ElementB);
    //乒乓缓冲：两个缓冲区轮流工作，一个负责接收数据，另一个交给CPU处理，互不干扰。

    // 申请L0A, L0B空间
    AscendC::TBuf<AscendC::TPosition::A2> l0aBuf;
    AscendC::TPipe l0aPipe;
    l0aPipe.InitBuffer(l0aBuf, ArchTag::L0ASize);
    AscendC::LocalTensor<uint8_t> l0aTensor = l0aBuf.Get<uint8_t>();
    l0aPipe.Destroy();

    AscendC::TBuf<AscendC::TPosition::B2> l0bBuf;
    AscendC::TPipe l0bPipe;
    l0bPipe.InitBuffer(l0bBuf, ArchTag::L0BSize);
    AscendC::LocalTensor<uint8_t> l0bTensor = l0bBuf.Get<uint8_t>();
    l0bPipe.Destroy();

    // 分配L0A, L0B空间
    AscendC::LocalTensor<ElementA> l0aTensorA[2];
    AscendC::LocalTensor<ElementB> l0bTensorB[2];


    l0aTensorA[0] = l0aTensor[0].template ReinterpretCast<ElementA>();
    l0aTensorA[1] = l0aTensor[L0A_PINGPONG_BUF_LEN * sizeof(ElementA)].template ReinterpretCast<ElementA>();

    l0bTensorB[0] = l0bTensor[0].template ReinterpretCast<ElementB>();
    l0bTensorB[1] = l0bTensor[L0B_PINGPONG_BUF_LEN * sizeof(ElementB)].template ReinterpretCast<ElementB>();

    // 计算 L0A, L0B能放下的数据大小
    uint32_t mL0 = mL1Round;
    uint32_t nL0 = nL1Round;
    uint32_t kL0 = min(
        RoundDown(L0A_PINGPONG_BUF_LEN / mL0, KL0AlignNumElement),
        RoundDown(L0B_PINGPONG_BUF_LEN / nL0, KL0AlignNumElement)
    );
    uint32_t kL0Loop = CeilDiv(kL1Round, kL0);
    for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
        uint32_t kL0Actual = (kL0Idx == kL0Loop - 1) ? (kL1Actual - kL0Idx * kL0) : kL0;
        uint32_t kL0Round = RoundUp(kL0Actual, KL0AlignNumElement);

        // 读取数据从L1到L0A
        l12l0_Nz2Zz<ElementA>(
            l0aTensorA[kL0Idx % 2],
            l1TensorA[kL0Idx * kL0 * mL0],
            mL0,
            kL0Round,
            mL1Round,
            ML0AlignNumElement,
            KL0AlignNumElement
        );

        // 读取数据从L1到L0B
        l12l0_Nz2Nz<ElementB>(
            l0bTensorB[kL0Idx % 2],
            l1TensorB[kL0Idx * kL0 * nL0],
            nL0,
            kL0Round,
            nL1Round,
            NL0AlignNumElement,
            KL0AlignNumElement
        );

        AscendC::PipeBarrier<PIPE_ALL>();

        // 计算矩阵乘法，输入在L0A,L0B;输出在L0C
        AscendC::Mmad(
            l0cTensorC,
            l0aTensorA[kL0Idx % 2],
            l0bTensorB[kL0Idx % 2],
            AscendC::MmadParams(
                mL0, 
                nL0, 
                kL0Actual, 
                0, 
                false, 
                (kL0Idx == 0) && isFirstLoop //我猜这玩意是控制是否累加的。
            )
        );

        AscendC::PipeBarrier<PIPE_ALL>();
    }
}

template<
    typename ArchTag,
    typename ElementA,
    typename ElementB,
    typename ComputeElementC
>
__aicore__ inline __attribute__((always_inline)) void matmulOnBlock_sync(
    AscendC::GlobalTensor<ElementA> gmTensorA,
    AscendC::GlobalTensor<ElementB> gmTensorB,
    AscendC::LocalTensor<ComputeElementC> l0cTensorC,
    uint32_t strideA,
    uint32_t strideB,
    uint32_t strideC,
    uint32_t maxMPerBlock,
    uint32_t maxNPerBlock,
    uint32_t maxKPerBlock,
    uint32_t MPerBlock,
    uint32_t NPerBlock,
    uint32_t KPerBlock
) {
    // 计算在L1上各方向应该对齐的元素个数
    uint32_t ML1AlignNumElement = C0_NUM_PER_FRACTAL;
    uint32_t NL1AlignNumElement = C0_NUM_PER_FRACTAL;
    uint32_t KL1AlignNumElement = BYTE_PER_C0 / sizeof(ElementA); 

    // 申请L1空间
    AscendC::TBuf<AscendC::TPosition::A1> l1Buf;
    AscendC::TPipe l1Pipe;
    l1Pipe.InitBuffer(l1Buf, ArchTag::L1Size);
    AscendC::LocalTensor<uint8_t> l1Tensor = l1Buf.Get<uint8_t>();
    l1Pipe.Destroy();

    // 分配L1空间给Tensor
    uint32_t L1A_SIZE = maxMPerBlock * maxKPerBlock;
    uint32_t L1B_SIZE = maxKPerBlock * maxNPerBlock;

    AscendC::LocalTensor<ElementA> l1TensorA[2]; //因为每次只是需要用一个Tensor，所以算上交替tensor在内，只需要两个即可
    AscendC::LocalTensor<ElementB> l1TensorB[2];

    uint32_t nowL1Bytes = 0;

    l1TensorA[0] = l1Tensor[nowL1Bytes].template ReinterpretCast<ElementA>();
    nowL1Bytes += L1A_SIZE * sizeof(ElementA);

    l1TensorA[1] = l1Tensor[nowL1Bytes].template ReinterpretCast<ElementA>();
    nowL1Bytes += L1A_SIZE * sizeof(ElementA);

    l1TensorB[0] = l1Tensor[nowL1Bytes].template ReinterpretCast<ElementB>();
    nowL1Bytes += L1B_SIZE * sizeof(ElementB);

    l1TensorB[1] = l1Tensor[nowL1Bytes].template ReinterpretCast<ElementB>();
    nowL1Bytes += L1B_SIZE * sizeof(ElementB);
    //这里实际上是把l1Tensor分成四份了，A、B各拿两个缓冲区。

    uint32_t mL1Actual = MPerBlock;
    uint32_t nL1Actual = NPerBlock;
    uint32_t mL1Round = RoundUp(mL1Actual, ML1AlignNumElement);
    uint32_t nL1Round = RoundUp(nL1Actual, NL1AlignNumElement);
    uint32_t kL1Loop = CeilDiv(KPerBlock, maxKPerBlock);//这里是按tile来的，把长度为K（也就是KPerBlock）的A的一行和B 的一列分成长度为maxKPerBlock的kL1Loop个。
    for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) {
        uint32_t kL1Actual = (kL1Idx == kL1Loop - 1 ? (KPerBlock - kL1Idx * maxKPerBlock) : maxKPerBlock); //判断不整除的部分
        uint32_t kL1Round = RoundUp(kL1Actual, KL1AlignNumElement);

        // 将数据从gm读取到L1
        gm2l1_Nd2Nz(
            l1TensorA[kL1Idx % 2], //轮着用L1A的两个cache
            gmTensorA[kL1Idx * maxKPerBlock], 
            //gm中每次取用maxKPerBlock长度的一个tile。
            //这里为什么可以直接从gmTensorA对应块的起始位置直接连续取数据？不是应该取m*k的一个小块吗？
            //我猜测，也许是因为这个矩阵会被以NZ的形式copy过来。它在datacopy的时候使用了一个Nd2NzParams参数。
            //一开始我还不理解为什么要以奇怪的Z字型存储矩阵，现在懂了。它会在读取小tile块的时候能更好的利用cache之类的东西。
            mL1Actual,
            kL1Actual,
            strideA,
            mL1Round,
            kL1Round
        );

        gm2l1_Nd2Nz(
            l1TensorB[kL1Idx % 2],
            gmTensorB[kL1Idx * maxKPerBlock],
            nL1Actual,
            kL1Actual,
            strideB,
            nL1Round,
            kL1Round
        );

        AscendC::PipeBarrier<PIPE_ALL>(); //这里是阻塞，让数据正常读取完
        
        // 计算矩阵乘法，输入在L1，输出在L0C
        matmulOnL1_sync<ArchTag, ElementA, ElementB, ComputeElementC>(
            l1TensorA[kL1Idx % 2],
            l1TensorB[kL1Idx % 2],
            l0cTensorC,
            mL1Actual,
            nL1Actual,
            kL1Actual,
            mL1Round,
            nL1Round,
            kL1Round,
            (kL1Idx == 0)
        );//把当前tile的结果输出累加在L0C中了。

        AscendC::PipeBarrier<PIPE_ALL>();
    }
}

template<
    typename ArchTag,
    typename ElementA,
    typename ElementB,
    typename ElementC,
    typename ComputeElementC,
    uint32_t maxMPerBlock,
    uint32_t maxNPerBlock,
    uint32_t maxKPerBlock
>
__aicore__ inline __attribute__((always_inline)) void matadd(
    __gm__ ElementA * ptr_A,
    __gm__ ElementB * ptr_B,
    __gm__ ElementC * ptr_C,
    uint32_t M,
    uint32_t N,
    uint32_t strideA,
    uint32_t strideB,
    uint32_t strideC
) {
    // uint32_t ML1AlignNumElement = C0_NUM_PER_FRACTAL;
    // uint32_t NL1AlignNumElement = C0_NUM_PER_FRACTAL;
    // //必须要做好内存对齐。C0_NUM_PER_FRACTAL是硬件的一个最小计算单元大小。数据维度必须对齐到这个粒度，硬件才能高效处理
    // //比如，对于MBlockRound = RoundUp(MBlockActual, ML1AlignNumElement);
    // //如果MBlockActual是100，而ML1AlignNumElement是16的话，MBlockRound会被对齐到128



    // // 申请gmTensor。把传入的数据复制到gmTensor里来了。
    // AscendC::GlobalTensor<ElementA> gmTensorA;
    // AscendC::GlobalTensor<ElementB> gmTensorB;
    // AscendC::GlobalTensor<ElementC> gmTensorC;
    // gmTensorA.SetGlobalBuffer((__gm__ ElementA *)ptr_A);
    // gmTensorB.SetGlobalBuffer((__gm__ ElementB *)ptr_B);
    // gmTensorC.SetGlobalBuffer((__gm__ ElementC *)ptr_C);

    // // 申请l0c。也就是存放结果矩阵C的L0cache部分的内容。在这声明，传入接下来的一系列计算中使得能够把结果传回。
    // //这里声明的比较复杂，也不需要理解为啥这么做，以后也跟着搞就好了
    // AscendC::TQue<AscendC::QuePosition::CO1,1>l0cTQue;
    // AscendC::TPipe l0cPipe;
    // l0cPipe.InitBuffer(l0cTQue, 1, ArchTag::L0CSize);
    // AscendC::LocalTensor<uint8_t> l0cTensor = l0cTQue.AllocTensor<uint8_t>();
    // l0cPipe.Destroy(); 

    // // 单个矩阵的大小
    // static constexpr uint32_t singleBlockSize = maxMPerBlock * maxNPerBlock;
    // // l0c能放矩阵的大小
    // //fp16中，ComputeElementC是float。
    // //ArchTag::L0CSize / sizeof(ComputeElementC)是能装多少个元素。再除以一个矩阵有多少个这个元素singleBlockSize即可得到结果矩阵的块数
    // static constexpr uint32_t l0c_blockNum = ArchTag::L0CSize / sizeof(ComputeElementC) / singleBlockSize;

    // // 申请存放多个矩阵的TensorC
    // AscendC::LocalTensor<ComputeElementC> l0cTensorC[l0c_blockNum];
    // #pragma unroll // 编译器指令：循环展开优化
    // for (uint32_t i = 0; i < l0c_blockNum; i++) {
    //     l0cTensorC[i] = l0cTensor[i * singleBlockSize * sizeof(ComputeElementC)].template ReinterpretCast<ComputeElementC>();
    //     //把uint8_t的l0cTensor换成每sizeof(ComputeElementC)个一个单元的l0cTensorC。
    //     //.template ReinterpretCast<ComputeElementC>()：类型转换，告诉编译器"把这块内存当作ComputeElementC类型来处理。因为上面声明的时候l0cTensor是uint8_t的tensor
    //     //为什么不直接AscendC::LocalTensor<ComputeElementC> l0cTensor = l0cTQue.AllocTensor<ComputeElementC>();呢？我搞不清楚。但是他这么写了我最好模仿他。
    // }

    // uint32_t MLoops = CeilDiv(M, maxMPerBlock);
    // uint32_t NLoops = CeilDiv(N, maxNPerBlock);
    // uint32_t coreLoops = MLoops * NLoops;
    // uint32_t singleLoopIdx = 0;
    // // 分核：将任务平均分给不同的核
    // for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) {
    //     /*
    //     分法：也还是把矩阵当成一维的铺开来看了。把C矩阵分成coreLoops个singleBlockSize大小（注意这里可能不能完整整除，会有些边角料块比这个大小小。
    //     然后块间采用经典的间隔BlockNum的方法取任务。
    //     */
    //     uint32_t MBlockIdx = loopIdx / NLoops; //当前块在第几行
    //     uint32_t NBlockIdx = loopIdx % NLoops; //当前块在第几列
    //     uint32_t MBlockActual = (MBlockIdx == MLoops - 1) ? (M - MBlockIdx * maxMPerBlock) : maxMPerBlock; //尾部不能整除maxMPerBlock的部分的处理
    //     uint32_t NBlockActual = (NBlockIdx == NLoops - 1) ? (N - NBlockIdx * maxNPerBlock) : maxNPerBlock;
    //     uint32_t MBlockRound = RoundUp(MBlockActual, ML1AlignNumElement);//如前面所说的，直接对齐最小的计算单元。
    //     uint32_t NBlockRound = RoundUp(NBlockActual, NL1AlignNumElement);
    //     // 计算(m, n, k) = (maxMPerBlock, maxNPerBlock, K) 的矩阵乘法，结果放在l0c中
    //     matmulOnBlock_sync<ArchTag, ElementA, ElementB, ComputeElementC>(
    //         gmTensorA[MBlockIdx * maxMPerBlock * strideA],
    //         //看看这个索引。我们需要的是计算C矩阵的(M,N)块。即需要A的第M行（以块为单位）。这里一块里面有maxMPerBlock行，有strideA列（这里的stride指的其实是块的列数）
    //         //因此前面的一个块有maxMPerBlock * strideA个元素。再乘上MBlockIdx就可以得到gmTensorA对应开始位置的下标了。需要的是
    //         gmTensorB[NBlockIdx * maxNPerBlock * strideB],
    //         //同A的索引，但是我们知道B是需要一列一列来的。所以这里B可能是列主次序的存储。
    //         l0cTensorC[singleLoopIdx % l0c_blockNum], //这里甚至循环使用l0cTensorC。因为L0CSize实际上是有限的。l0c_blockNum是我们一共能使用的Loc块数，而不是完成任务需要的块数
    //         strideA,
    //         strideB,
    //         strideC,
    //         maxMPerBlock,
    //         maxNPerBlock,
    //         maxKPerBlock,
    //         MBlockActual,
    //         NBlockActual,
    //         K
    //     );

    //     AscendC::PipeBarrier<PIPE_ALL>();

    //     // 将l0c的结果写入到全局内存。我们的C为了方便是NZ的形式存的，这里Nz2Nd转回Nd
    //     l0c2gm_Nz2Nd<ElementC, ComputeElementC>(
    //         gmTensorC[MBlockIdx * maxMPerBlock * strideC + NBlockIdx * maxNPerBlock],//这个就是经典的找到(M, N)对应的块的元素下标了
    //         l0cTensorC[singleLoopIdx % l0c_blockNum],
    //         MBlockActual,
    //         NBlockActual,
    //         MBlockRound,
    //         NBlockRound,
    //         strideC
    //     );

    //     AscendC::PipeBarrier<PIPE_ALL>();

    //     singleLoopIdx++;

    // }
}